#!/usr/bin/env python3
"""
DataPlexus 修复 batch_size NULL 值脚本
修复 dp_sync_tasks 表中 batch_size 为 NULL 的记录
"""

import pymysql
import sys
import os

# 数据库配置
DB_CONFIG = {
    'host': '192.168.37.11',
    'port': 3306,
    'user': 'root',
    'password': '123456',
    'database': 'data_plexus',
    'charset': 'utf8mb4'
}

def check_and_fix_batch_size():
    """检查并修复 batch_size NULL 值"""
    connection = None
    try:
        # 连接数据库
        print("连接数据库...")
        connection = pymysql.connect(**DB_CONFIG)
        cursor = connection.cursor()
        
        print(f"成功连接到数据库: {DB_CONFIG['host']}:{DB_CONFIG['port']}/{DB_CONFIG['database']}")
        
        # 检查表是否存在
        cursor.execute("""
            SELECT COUNT(*) 
            FROM INFORMATION_SCHEMA.TABLES 
            WHERE TABLE_SCHEMA = %s 
              AND TABLE_NAME = 'dp_sync_tasks'
        """, (DB_CONFIG['database'],))
        
        if cursor.fetchone()[0] == 0:
            print("错误: dp_sync_tasks表不存在")
            return False
        
        # 检查当前 batch_size 为 NULL 的记录
        print("\n检查当前 batch_size 为 NULL 的记录...")
        cursor.execute("""
            SELECT 
                id,
                name,
                display_name,
                sync_type,
                batch_size,
                is_active
            FROM dp_sync_tasks 
            WHERE batch_size IS NULL AND is_deleted = FALSE
            ORDER BY id
        """)
        
        null_records = cursor.fetchall()
        if null_records:
            print(f"发现 {len(null_records)} 个 batch_size 为 NULL 的记录:")
            print("ID\t名称\t\t\t显示名称\t\t同步类型\tbatch_size\t状态")
            print("-" * 80)
            for record in null_records:
                id, name, display_name, sync_type, batch_size, is_active = record
                status = "启用" if is_active else "禁用"
                print(f"{id}\t{name[:15]:<15}\t{display_name[:15]:<15}\t{sync_type}\t{batch_size}\t{status}")
        else:
            print("没有发现 batch_size 为 NULL 的记录")
        
        # 显示统计信息
        print("\n当前统计信息:")
        cursor.execute("""
            SELECT 
                COUNT(*) as total_tasks,
                COUNT(CASE WHEN batch_size IS NULL THEN 1 END) as null_batch_size_count,
                COUNT(CASE WHEN batch_size IS NOT NULL THEN 1 END) as valid_batch_size_count
            FROM dp_sync_tasks 
            WHERE is_deleted = FALSE
        """)
        
        stats = cursor.fetchone()
        total_tasks, null_count, valid_count = stats
        print(f"总任务数: {total_tasks}")
        print(f"batch_size 为 NULL 的任务数: {null_count}")
        print(f"batch_size 有效的任务数: {valid_count}")
        
        if null_count == 0:
            print("\n✓ 所有任务的 batch_size 都已设置，无需修复")
            return True
        
        # 询问是否执行修复
        print(f"\n是否要修复这 {null_count} 个记录的 batch_size 值？(y/N): ", end="")
        response = input().strip().lower()
        
        if response not in ['y', 'yes']:
            print("取消修复操作")
            return True
        
        # 执行修复
        print("\n开始修复 batch_size NULL 值...")
        cursor.execute("""
            UPDATE dp_sync_tasks 
            SET 
                batch_size = 1000,
                updated_at = CURRENT_TIMESTAMP
            WHERE batch_size IS NULL 
              AND is_deleted = FALSE
        """)
        
        updated_rows = cursor.rowcount
        connection.commit()
        
        print(f"✓ 成功更新了 {updated_rows} 个记录的 batch_size 值")
        
        # 验证修复结果
        print("\n验证修复结果...")
        cursor.execute("""
            SELECT 
                COUNT(*) as total_tasks,
                COUNT(CASE WHEN batch_size IS NULL THEN 1 END) as null_batch_size_count,
                COUNT(CASE WHEN batch_size IS NOT NULL THEN 1 END) as valid_batch_size_count,
                MIN(batch_size) as min_batch_size,
                MAX(batch_size) as max_batch_size,
                AVG(batch_size) as avg_batch_size
            FROM dp_sync_tasks 
            WHERE is_deleted = FALSE
        """)
        
        final_stats = cursor.fetchone()
        total, null_final, valid_final, min_batch, max_batch, avg_batch = final_stats
        
        print(f"修复后统计:")
        print(f"总任务数: {total}")
        print(f"batch_size 为 NULL 的任务数: {null_final}")
        print(f"batch_size 有效的任务数: {valid_final}")
        print(f"batch_size 最小值: {min_batch}")
        print(f"batch_size 最大值: {max_batch}")
        print(f"batch_size 平均值: {avg_batch:.2f}")
        
        if null_final == 0:
            print("\n✓ 所有 batch_size NULL 值已成功修复!")
        else:
            print(f"\n⚠ 仍有 {null_final} 个记录的 batch_size 为 NULL")
        
        return null_final == 0
        
    except Exception as e:
        print(f"修复失败: {e}")
        if connection:
            connection.rollback()
        return False
        
    finally:
        if connection:
            connection.close()
            print("\n数据库连接已关闭")

def main():
    """主函数"""
    print("DataPlexus batch_size NULL 值修复工具")
    print("=" * 50)
    
    success = check_and_fix_batch_size()
    
    if success:
        print("\n✓ 修复完成!")
        sys.exit(0)
    else:
        print("\n✗ 修复失败!")
        sys.exit(1)

if __name__ == "__main__":
    main()
